package com.feidee.transformer;

import com.feidee.fdspark.transformer.WordSpliter;
import org.dmg.pmml.DataType;
import org.dmg.pmml.FieldName;
import org.dmg.pmml.OpType;
import org.jpmml.converter.ContinuousFeature;
import org.jpmml.converter.Feature;
import org.jpmml.converter.FeatureUtil;
import org.jpmml.converter.PMMLUtil;
import org.jpmml.sparkml.DocumentFeature;
import org.jpmml.sparkml.FeatureConverter;
import org.jpmml.sparkml.SparkMLEncoder;

import java.util.Collections;
import java.util.List;

/**
 * @Author songhaicheng
 * @Date 2018/12/19 15:23
 * @Description
 * @Reviewer
 */
public class WordSpliterConverter extends FeatureConverter<WordSpliter> {

    public static final String DEFAULT_DELIMITER = " ";

    public WordSpliterConverter(WordSpliter transformer) {
        super(transformer);
    }

    @Override
    public List<Feature> encodeFeatures(SparkMLEncoder encoder) {
        Feature inputFeature = encoder.getOnlyFeature(getTransformer().getInputCol());
        // 输出结果列
        switch (getTransformer().getToType().toString()) {
            case "Array": {
                return Collections.singletonList(new DocumentFeature(
                        encoder,
                        encoder.createDerivedField(
                                FeatureUtil.createName("lowercase", inputFeature),
                                OpType.CATEGORICAL,
                                DataType.STRING,
                                PMMLUtil.createApply("lowercase", inputFeature.ref())),
                                "\\s+")
                );
            }
            default: {      // String
                // 替换原分隔符成 ","
                if (!DEFAULT_DELIMITER.equals(getTransformer().getDelimiter())) {
                    return Collections.singletonList(new ContinuousFeature(
                            encoder,
                            encoder.createDerivedField(
                                    FieldName.create(getTransformer().getOutputCol()),
                                    OpType.CATEGORICAL,
                                    DataType.STRING,
                                    PMMLUtil.createApply(
                                            "replace",
                                            inputFeature.ref(),
                                            PMMLUtil.createConstant(getTransformer().getDelimiter(), DataType.STRING),
                                            PMMLUtil.createConstant(DEFAULT_DELIMITER, DataType.STRING)
                                    )
                            ))
                    );
                } else {
                    // 用空格分隔时，不做任何处理
                    break;
                }
            }
        }
        return Collections.singletonList(inputFeature);
    }
}
